package dogapm

import (
	"context"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
	"net"
	"time"
)

var grpcServerTracerName = "grpc-server"

type GrpcServer struct {
	*grpc.Server
	addr string
}

func NewGrpcServer(addr string) *GrpcServer {
	server := &GrpcServer{
		Server: grpc.NewServer(grpc.UnaryInterceptor(unaryServerInterceptor())),
		addr:   addr,
	}
	globalStarters = append(globalStarters, server)
	globalClosers = append(globalClosers, server)
	return server
}
func (g *GrpcServer) Close() {
	g.Server.GracefulStop()
}

func unaryServerInterceptor() grpc.UnaryServerInterceptor {
	tracer := otel.GetTracerProvider().Tracer(grpcServerTracerName)
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
		md, ok := metadata.FromIncomingContext(ctx)
		if !ok {
			md = metadata.MD{}
		}

		// 将md 注入到ctx
		ctx = otel.GetTextMapPropagator().Extract(ctx, &metadataSupplier{
			metadata: &md,
		})

		ctx, span := tracer.Start(ctx, info.FullMethod, trace.WithSpanKind(trace.SpanKindServer))
		start := time.Now()
		defer func() {
			span.SetAttributes(attribute.Int64("grpc.duration", time.Now().Sub(start).Milliseconds()))
			span.End()
		}()
		resp, err = handler(ctx, req)
		if err != nil {
			s, _ := status.FromError(err)
			span.RecordError(err, trace.WithStackTrace(true), trace.WithTimestamp(time.Now()))
			span.SetAttributes(attribute.Bool("error", true))
			span.SetAttributes(attribute.String("grpc.status_code", s.Code().String()))
			return resp, err
		}
		return resp, nil
	}
}
func (g *GrpcServer) Start() {
	l, err := net.Listen("tcp", g.addr)
	if err != nil {
		panic(err)
	}
	go func() {
		err = g.Serve(l)
		if err != nil {
			panic(err)
		}
	}()
}
